import tkinter as tk
import messagebox
import cv2
import pyautogui
import numpy

try:
    import tkinter as tk
    from tkinter import messagebox
except ImportError:
    print("无法导入 tkinter 库，请检查 Python 环境是否支持。")
    raise

try:
    import cv2
    import pyautogui
    import numpy as np
    from ultralytics import YOLO
    import win32gui
    import win32con
except ImportError as e:
    missing_libraries = str(e).split("'")[1]

    def show_error_and_exit():
        root = tk.Tk()
        root.title("缺少库")

        # 获取屏幕尺寸
        screen_width = root.winfo_screenwidth()
        screen_height = root.winfo_screenheight()

        window_width = 300
        window_height = 150

        # 计算窗口在屏幕中心的坐标
        x = (screen_width - window_width) // 2
        y = (screen_height - window_height) // 2

        root.geometry(f"{window_width}x{window_height}+{x}+{y}")

        label = tk.Label(root, text=f"程序无法运行，缺少以下库: {missing_libraries}，请使用 pip install {missing_libraries} 进行安装。")
        label.pack(pady=20)

        def on_ok_click():
            root.destroy()
            raise SystemExit

        ok_button = tk.Button(root, text="确定", command=on_ok_click)
        ok_button.pack(pady=20)

        root.mainloop()

    show_error_and_exit()


# 启动窗口线程
cv2.startWindowThread()

# 加载本地 YOLOv8 模型
model = YOLO('ldjsqV1/detection/best.pt')  # 替换为你的模型路径

# 创建窗口并设置大小
cv2.namedWindow('Detection Result', cv2.WINDOW_NORMAL)
cv2.resizeWindow('Detection Result', 300, 300)

hwnd = win32gui.FindWindow(None, 'Detection Result')
if hwnd == 0:
    print("未找到窗口句柄，请检查窗口名称。")
else:
    try:
        win32gui.SetWindowPos(hwnd, win32con.HWND_TOPMOST, 0, 0, 0, 0,
                              win32con.SWP_NOMOVE | win32con.SWP_NOSIZE)
        print("窗口已置顶。")
    except Exception as e:
        print(f"窗口置顶失败: {e}")

while True:
    # 获取屏幕尺寸
    screen_width, screen_height = pyautogui.size()

    # 计算屏幕中心
    center_x = screen_width // 2
    center_y = screen_height // 2

    # 计算 320x320 区域的左上角坐标
    x1 = center_x - 160
    y1 = center_y - 160

    try:
        # 截取屏幕中心 320x320 大小的图像
        screenshot = pyautogui.screenshot(region=(x1, y1, 320, 320))
        frame = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
    except Exception as e:
        print(f"屏幕截图失败: {e}")
        continue

    # 使用 YOLOv8 模型进行检测
    results = model(frame)

    # 获取检测结果图像
    result_img = results[0].plot()

    # 调整检测结果图像大小为 300x300
    resized_result = cv2.resize(result_img, (300, 300))

    # 显示图像
    cv2.imshow('Detection Result', resized_result)

    # 按 'q' 键退出循环
    key = cv2.waitKey(1)
    if key & 0xFF == ord('q'):
        break

    # 检查窗口是否被关闭
    if cv2.getWindowProperty('Detection Result', cv2.WND_PROP_VISIBLE) < 1:
        break

# 关闭所有窗口
cv2.destroyAllWindows()